import cv2
import os
import numpy as np
from ultralytics import YOLO
import shutil

def create_directories():
    """创建必要的目录结构"""
    dirs = ['dataset/train', 'dataset/val', 'dataset/test']
    for dir_path in dirs:
        if os.path.exists(dir_path):
            shutil.rmtree(dir_path)
        os.makedirs(dir_path)

def extract_frames(video_path, save_interval=10):
    """
    从视频中提取帧
    video_path: 视频文件路径
    save_interval: 每隔多少帧保存一帧
    """
    # 加载YOLOv8-pose模型
    model = YOLO('yolov8n-pose.pt')
    
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    frame_count = 0
    saved_frames = []
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
            
        if frame_count % save_interval == 0:
            # 使用YOLOv8进行姿态检测
            results = model(frame)
            
            # 如果检测到人体关键点，保存该帧
            if len(results[0].keypoints.data) > 0:
                saved_frames.append(frame)
                
        frame_count += 1
        
    cap.release()
    return saved_frames

def split_and_save_frames(frames):
    """
    将帧分割为训练集、验证集和测试集
    """
    total_frames = len(frames)
    train_size = int(0.8 * total_frames)
    val_size = int(0.1 * total_frames)
    
    # 随机打乱帧的顺序
    indices = np.random.permutation(total_frames)
    
    # 分割数据集
    train_indices = indices[:train_size]
    val_indices = indices[train_size:train_size+val_size]
    test_indices = indices[train_size+val_size:]
    
    # 保存帧到相应目录
    for idx in train_indices:
        cv2.imwrite(f'dataset/train/frame_{idx:04d}.jpg', frames[idx])
    
    for idx in val_indices:
        cv2.imwrite(f'dataset/val/frame_{idx:04d}.jpg', frames[idx])
        
    for idx in test_indices:
        cv2.imwrite(f'dataset/test/frame_{idx:04d}.jpg', frames[idx])

def main():
    create_directories()
    frames = extract_frames('match.mp4')
    split_and_save_frames(frames)
    print(f"总共处理的帧数: {len(frames)}")
    print(f"训练集帧数: {int(0.8 * len(frames))}")
    print(f"验证集帧数: {int(0.1 * len(frames))}")
    print(f"测试集帧数: {int(0.1 * len(frames))}")

if __name__ == "__main__":
    main() 